#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
#
# This file is part of Coruja-scripts
#
# Coruja-scripts is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 3 of the License.
#
# Coruja-scripts is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Coruja-scripts.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright 2011 Grupo Falabrasil - http://www.laps.ufpa.br/falabrasil
#
# Author 2011: Pedro Batista - pedosb@gmail.com

import argparse
import sys
import random
import re
import numpy as np
import vq
from config import *
try:
	from energy_estimation import get_htk_parameters
except ImportError:
	pass
from dtw import dtw_distance
from HTKParameters import HTKParameters as hpm

from mpi4py import MPI
comm = MPI.COMM_WORLD
size = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
name = MPI.Get_processor_name()

def decode_vq(word_list=config.get('files', 'word_list'),
		test_file_list=config.get('files', 'test_file_list'),
		template=config.get('files', 'template'),
		word_from_feat_file_regex=config.get('files', 'word_from_feat_file_regex'),
		speaker_from_feat_file_regex=config.get('files',
			'speaker_from_feat_file_regex'),
		result_file=config.get('files', 'result_file')):
	cb = vq.Codebook(
			code_word_list=[hpm(t.strip()) for t in open(template).readlines()])
	word_from_feat = re.compile(word_from_feat_file_regex).match
	speaker_from_feat = re.compile(speaker_from_feat_file_regex).match
	if rank == 0:
		l = [l.strip() for l in open(test_file_list).readlines()]
		random.shuffle(l)
	else:
		l = None
	l = comm.bcast(l, root=0)
	m = len(l) / size
	resto = len(l) - m * size
	s = rank * m + (rank if rank < resto else resto)
	l = l[s : s + m + (1 if rank < resto else 0)]
	r = []
	speaker_err = dict()
	speaker_acc = dict()
	for t in l:
		ar = (word_from_feat(t).group(1), word_from_feat(
			cb.code_word_list[cb.classify(hpm(t))]._file.name).group(1))
		r.append(ar)
		speaker = speaker_from_feat(t).group(1)
		if ar[0] != ar[1]:
			try:
				speaker_err[speaker] += 1
			except KeyError:
				speaker_err[speaker] = 1
		else:
			try:
				speaker_acc[speaker] += 1
			except KeyError:
				speaker_acc[speaker] = 1
	if rank == 0:
		out = open(result_file, 'w')
		out.write('Label\tRecognized\n')
		out.write('\n'.join(['\t'.join(i) for i in r]) + '\n')
		out.flush()
		for i in xrange(1, size):
			r = comm.recv(source=i, tag=i)
			speaker_err = concat_dict(speaker_err,
					comm.recv(source=i, tag=i+size+1))
			speaker_acc = concat_dict(speaker_acc,
					comm.recv(source=i, tag=i+size+1))
			out.write('\n'.join(['\t'.join(i) for i in r]) + '\n')
			out.flush()
		out.close()
		report(speaker_acc, speaker_err)
	else:
		comm.send(r, dest=0, tag=rank)
		comm.send(speaker_err, dest=0, tag=rank+size+1)
		comm.send(speaker_acc, dest=0, tag=rank+size+1)

def concat_dict(*dict_list):
	nd = dict()
	for d in dict_list:
		for k in d:
			try:
				nd[k] += d[k]
			except KeyError:
				nd[k] = d[k]
	return nd

def report(acc, err):
	for speaker in set(acc.keys() + err.keys()):
		if (speaker in err.keys() and 
				speaker in acc.keys()):
			error = float(err[speaker]) / (acc[speaker] + err[speaker])
			files = err[speaker] + acc[speaker]
		elif speaker in acc.keys():
			error = 0
			files = acc[speaker]
		else:
			error = 1
			files = err[speaker]
		print '%s:\t %.2f of %d files' % (speaker, error * 100,
				files)

def decode(test_file, template_list):
	if isinstance(template_list, str):
		with open(template_list) as f:
			template_list = [l.strip() for l in f.readlines()]
	distances = []
	size = 100 / float(len(template_list))
	i = 0
	for t in template_list:
		i += 1
		sys.stderr.write('\r%6.2f%%' % (i * size))
		sys.stderr.flush()
		if not isinstance(test_file, hpm):
			try:
				test_file = get_htk_parameters(test_file)
			except ValueError:
				pass
#		print t
		distances.append(dtw_distance(test_file, t))
	sys.stderr.write('\n')
	distances = np.array(distances)
	print template_list[distances.argmin()]

if __name__=='__main__':
	parser = argparse.ArgumentParser(description='Decode test with templates.')
#	parser.add_argument('-t', '--template-file-list', required=True,
#		help='File with a template (one by each line)')
#	parser.add_argument('-f', '--test', required=True,
#		help='Test file')
#	parser.add_argument('-c', '--config-file', default='',
#		help='A specific config file')
	args = parser.parse_args()
#	load_config(args.config_file)
#	decode(args.test, args.template_file_list)
	decode_vq()
